import torch
from sklearn import metrics

class MetricsUtil:
    @staticmethod
    def cal_metrics(predicts, labels):
        n_test_total = len(labels)
        n_test_correct = (torch.argmax(predicts, -1) == labels).sum().item()

        accuracy = n_test_correct / n_test_total
        f1_macro = metrics.f1_score(labels.cpu(), torch.argmax(predicts, -1).cpu(), labels=[0, 2], average='macro')
        f1_micro = metrics.f1_score(labels.cpu(), torch.argmax(predicts, -1).cpu(), labels=[0, 2], average='micro')
        f1_average = 0.5 * (f1_macro + f1_micro)

        return {f'accuracy':accuracy, f'f1_macro':f1_macro, f'f1_average':f1_average}